package org.snail.proxy.server.netty.ws;

import com.google.protobuf.ByteString;
import io.netty.channel.Channel;
import io.netty.channel.ChannelHandler;
import io.netty.channel.ChannelHandlerContext;
import io.netty.channel.SimpleChannelInboundHandler;
import io.netty.handler.codec.http.websocketx.BinaryWebSocketFrame;
import io.netty.util.Attribute;
import io.netty.util.AttributeKey;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.snail.proxy.proto.CreateClientRequest;
import org.snail.proxy.proto.DeleteClientRequest;
import org.snail.proxy.proto.MessageId;
import org.snail.proxy.proto.Packet;
import org.snail.proxy.proto.ProxyMessage;
import org.snail.proxy.server.netty.ProxyClientHandler;

import java.util.UUID;

@ChannelHandler.Sharable
public class WebsocketProxyHandler extends SimpleChannelInboundHandler<BinaryWebSocketFrame> {
    private static final Logger log = LoggerFactory.getLogger(WebsocketProxyHandler.class);

    private final AttributeKey<String> idKey = AttributeKey.valueOf("idKey");

    private final Channel proxyClientChannel;
    private final ProxyClientHandler proxyClientHandler;

    public WebsocketProxyHandler(Channel proxyClientChannel, ProxyClientHandler proxyClientHandler) {
        this.proxyClientChannel = proxyClientChannel;
        this.proxyClientHandler = proxyClientHandler;
    }

    @Override
    protected void channelRead0(ChannelHandlerContext ctx, BinaryWebSocketFrame msg) {
        Channel channel = ctx.channel();
        String id = channel.attr(idKey).get();
        if (id != null) {
            byte[] data = msg.content().array();
            ProxyMessage proxyMessage = ProxyMessage.newBuilder().setId(id).setData(ByteString.copyFrom(data)).build();
            proxyClientChannel.writeAndFlush(new Packet((byte) 0, MessageId.PROXY_MESSAGE_VALUE, proxyMessage.toByteArray()));
            log.info("Transform websocket message {}[{}].", id, data.length);
        }
    }

    @Override
    public void channelActive(ChannelHandlerContext ctx) throws Exception {
        Channel channel = ctx.channel();
        Attribute<String> attribute = channel.attr(idKey);
        if (attribute.get() == null) {
            String id = UUID.randomUUID().toString();
            attribute.set(id);
            proxyClientHandler.addProxyChannel(id, channel);
            CreateClientRequest request = CreateClientRequest.newBuilder().setId(id).build();
            proxyClientChannel.writeAndFlush(new Packet((byte) 0, MessageId.CREATE_CLIENT_REQUEST_VALUE, request.toByteArray()));
            log.info("Connect client {} from {}.", id, channel.remoteAddress());
        }
        super.channelActive(ctx);
    }

    @Override
    public void channelInactive(ChannelHandlerContext ctx) throws Exception {
        Channel channel = ctx.channel();
        String id = channel.attr(idKey).get();
        if (id != null) {
            proxyClientHandler.removeProxyChannel(id);
            DeleteClientRequest request = DeleteClientRequest.newBuilder().setId(id).build();
            proxyClientChannel.writeAndFlush(new Packet((byte) 0, MessageId.DELETE_CLIENT_REQUEST_VALUE, request.toByteArray()));
            log.info("Disconnect client {} from {}.", id, channel.remoteAddress());
        }
        super.channelInactive(ctx);
    }
}
